package zhipu

import (
	"context"
	"encoding/json"
	"errors"
	"log/slog"
	"quipus/internal/llm"

	mark3labs "github.com/mark3labs/mcp-go/mcp"
	zp "github.com/yankeguo/zhipu"
)

type ZhiPuClient struct {
	client *zp.Client
	model  string
}

type ZhiPuAuthInfo struct {
	SecretKey string `json:"secret_key"`
}

func NewClient(config *llm.LLMClientConfig) (*ZhiPuClient, error) {
	authInfo := ZhiPuAuthInfo{}
	err := json.Unmarshal([]byte(config.AuthInfo), &authInfo)
	if err != nil {
		return nil, err
	}
	client, err := zp.NewClient(zp.WithAPIKey(authInfo.SecretKey), zp.WithBaseURL(config.BaseUrl))

	if err != nil {
		return nil, err
	}

	c := &ZhiPuClient{
		client: client,
		model:  config.Model,
	}
	return c, nil
}

func (c *ZhiPuClient) Chat(ctx context.Context, req *llm.ChatReq) error {
	if req == nil {
		return nil
	}

	model := req.Model
	msgs := req.Msgs

	if len(msgs) < 1 {
		return errors.New("message is null")
	}

	if len(model) < 1 {
		model = c.model
	}

	service := c.client.ChatCompletion(model)
	for _, m := range msgs {
		service = service.AddMessage(zp.ChatCompletionMessage{
			Role:    m.Role,
			Content: m.Content,
		})
	}

	respFunc := func(chunk zp.ChatCompletionResponse) error {
		choices := chunk.Choices
		for _, choice := range choices {
			select {
			case req.RspChan <- choice.Delta.Content:
				// slog.Info("get response: ", choice.Delta.Content)
			case <-ctx.Done():
				slog.Info("reve quit signal ")

				return nil
			}
		}
		return nil
	}

	service.SetStreamHandler(respFunc)

	_, err := service.Do(ctx)
	if err != nil {
		req.CompleteChan <- true
		return err
	}

	// s := rsp.Choices[0].Message.Content
	// fmt.Print(s)
	req.CompleteChan <- true
	return nil
}

func (c *ZhiPuClient) GetProviderName() string {
	return "zhipu"
}

func (c *ZhiPuClient) Embedding(ctx context.Context, req *llm.EmbeddingReq) ([]float64, error) {
	service := c.client.Embedding(req.Model).SetInput(req.Prompt)
	response, err := service.Do(context.Background())
	if err != nil {
		return nil, err
	}

	return response.Data[0].Embedding, nil
}

func (c *ZhiPuClient) Complete(ctx context.Context, req *llm.CompleteReq) (llm.Message, error) {
	if req == nil {
		return nil, nil
	}

	model := req.Model
	msgs := req.Msgs

	if len(msgs) < 1 {
		return nil, errors.New("message is null")
	}

	if len(model) < 1 {
		model = c.model
	}

	service := c.client.ChatCompletion(model)
	for _, m := range msgs {
		service = service.AddMessage(zp.ChatCompletionMessage{
			Role:    m.Role,
			Content: m.Content,
		})
	}

	if req.Tools != nil && len(req.Tools) > 0 {
		service.AddTool(convertTool(req.Tools)...)
	}

	resp, err := service.Do(ctx)
	if err != nil {

		return nil, err
	}

	return &ZhiPuMessage{Resp: &resp, Choice: &resp.Choices[0]}, nil
}

func convertTool(tools []mark3labs.Tool) []zp.ChatCompletionTool {
	var zpTools []zp.ChatCompletionTool
	for _, tool := range tools {
		t := zp.ChatCompletionToolFunction{
			Name:        tool.Name,
			Description: tool.Description,
			Parameters:  convertSchema(tool.InputSchema),
		}
		zpTools = append(zpTools, t)
	}
	return zpTools
}
func convertSchema(schema mark3labs.ToolInputSchema) map[string]interface{} {
	required := schema.Required
	if required == nil {
		required = []string{}
	}

	return map[string]interface{}{
		"type":       schema.Type,
		"properties": schema.Properties,
		"required":   required,
	}
}

type ZhiPuMessage struct {
	Resp   *zp.ChatCompletionResponse
	Choice *zp.ChatCompletionChoice
}

func (m *ZhiPuMessage) GetRole() string {
	return string(m.Choice.Message.Role)
}

func (m *ZhiPuMessage) GetContent() string {
	return m.Choice.Message.Content
}

func (m *ZhiPuMessage) GetToolCalls() []llm.ToolCall {
	var calls []llm.ToolCall
	for _, call := range m.Choice.Message.ToolCalls {
		calls = append(calls, &ToolCallWrapper{call})
	}
	return calls
}

func (m *ZhiPuMessage) IsToolResponse() bool {
	return m.Choice.Message.ToolCalls != nil && len(m.Choice.Message.ToolCalls) > 0
}

func (m *ZhiPuMessage) GetToolResponseID() string {
	return m.Choice.Message.ToolCalls[0].ID
}

func (m *ZhiPuMessage) GetUsage() (int, int) {
	return int(m.Resp.Usage.CompletionTokens), int(m.Resp.Usage.PromptTokens)
}

type ToolCallWrapper struct {
	Call zp.ChatCompletionToolCall
}

func (t *ToolCallWrapper) GetID() string {
	return t.Call.ID
}

func (t *ToolCallWrapper) GetName() string {
	return t.Call.Function.Name
}

func (t *ToolCallWrapper) GetArguments() map[string]interface{} {
	var args map[string]interface{}
	if err := json.Unmarshal([]byte(t.Call.Function.Arguments), &args); err != nil {
		return make(map[string]interface{})
	}
	return args
}
